import streamlit as st
from scipy.optimize import linprog
import numpy as np
from database.queries.user_orm import SynOrmUserQueries
from widgets.chart_pfc import get_pfc_pie_figure


def __change_daily_norm(default_norm: tuple, key_base: str):
    if "username" not in st.session_state:
        return
    user = SynOrmUserQueries.get_user_by_username(st.session_state["username"])
    
    # Norm changing by user
    if f"{key_base}_daily_norm" not in st.session_state:
        st.session_state[f"{key_base}_daily_norm"] = user.get_daily_norm()
    
    if st.checkbox("Поменять распределение?", key=f"{key_base}_change_available"):
        def change_norm():
            st.session_state[f"{key_base}_daily_norm"] = (
                st.session_state[f"{key_base}_protein_change"], 
                st.session_state[f"{key_base}_fat_change"], 
                st.session_state[f"{key_base}_carb_change"], 
            ) 
            
            
        dn = st.session_state[f"{key_base}_daily_norm"]
        with st.form(f"{key_base}_change_daily_norm_form"):
            prot_col, fat_col, carb_col = st.columns(3)
            with prot_col:
                st.number_input(
                    label="Белки:", 
                    value=round(dn[0], 1),
                    min_value=0.0, 
                    step=0.1, 
                    key=f"{key_base}_protein_change",
                )
            with fat_col:
                st.number_input(
                    label="Жиры:", 
                    value=round(dn[1], 1), 
                    min_value=0.0, 
                    step=0.1, 
                    key=f"{key_base}_fat_change",
                )
            with carb_col:
                st.number_input(
                    label="Углеводы:", 
                    value=round(dn[2], 1), 
                    min_value=0.0, 
                    step=0.1, 
                    key=f"{key_base}_carb_change",
                )
            st.form_submit_button(
                label="Подтвердить", 
                type="primary",
                on_click=change_norm, 
            )
        st.warning(
            "**Не рекомендуется** менять распределение БЖУ, "
            "если вы не уверены в том, что делаете!",
        )
    else:
        st.session_state[f"{key_base}_daily_norm"] = default_norm
        
    return st.session_state[f"{key_base}_daily_norm"]
        
        
def __dislike_products(key_base: str):
    if "username" not in st.session_state:
        return
    user = SynOrmUserQueries.get_user_by_username(st.session_state["username"])
    
    # Take into account products that are disliked by user today
    if f"{key_base}_disliked_products" not in st.session_state:
        st.session_state[f"{key_base}_disliked_products"] = []
    disliked_prod_names = st.session_state[f"{key_base}_disliked_products"]
    for dpn in disliked_prod_names:
        # If removed during session
        if dpn not in [p.name for p in user.products]:
            disliked_prod_names.remove(dpn)
    
    st.session_state[f"{key_base}_available_products"] = sorted(
        [
            product.name for product in user.products 
            if product.name not in disliked_prod_names
        ]
    )
    available_prod_names = st.session_state[f"{key_base}_available_products"]
    
    with st.expander(label="❌ Не буду сегодня есть", expanded=True):
        def __dislike_product():
            product_name = st.session_state[
                f"{key_base}_select_product_to_dislike"
            ]
            if product_name is None:
                return
            if product_name not in disliked_prod_names:
                disliked_prod_names.append(product_name)
                disliked_prod_names.sort()
                available_prod_names.remove(product_name)
                available_prod_names.sort()
                liked_products = st.session_state[f"{key_base}_liked_products"]
                if liked_products is not None:
                    for product in liked_products:
                        if product.name == product_name:
                            liked_products.remove(product)
        
        with st.form(key=f"{key_base}_dislike_product_form"):
            st.selectbox(
                label="Что не хочу", 
                options=available_prod_names,
                key=f"{key_base}_select_product_to_dislike",
            )
            st.form_submit_button(
                label="Подтвердить исключение из рациона", 
                type="primary",
                on_click=__dislike_product,
            )
        
        def __return_disliked_product(product_name):
            disliked_prod_names.remove(product_name)
            disliked_prod_names.sort()
            available_prod_names.append(product_name)
            available_prod_names.sort()

        for i in range(len(disliked_prod_names)):
            name_col, return_col = st.columns(2)
            product = disliked_prod_names[i]
            with name_col:
                st.write(product)
            with return_col:
                st.button(
                    label="Вернуть", 
                    key=f"{key_base}_return_product_{i}_from_disliked", 
                    on_click=__return_disliked_product, 
                    args=[product],
                    use_container_width=True,
                )
    
    return disliked_prod_names


class __LikedProduct:
        def __init__(
            self,
            product_name: str, 
            min_weight: int, 
            max_weight: int,
        ):
            self.name = product_name
            self.min_weight = min_weight
            self.max_weight = max_weight


def __like_products(proteins_norm, fats_norm, carbs_norm, key_base: str):
    if "username" not in st.session_state:
        return
    user = SynOrmUserQueries.get_user_by_username(st.session_state["username"])
    
    # Take into account products that user wants to eat today
    if f"{key_base}_liked_products" not in st.session_state:
        st.session_state[f"{key_base}_liked_products"] = []
    liked_products = st.session_state[f"{key_base}_liked_products"]
    for lp in liked_products:
        # If removed during session
        if lp.name not in [p.name for p in user.products]:
            liked_products.remove(lp)
    
    available_products = None
    if f"{key_base}_available_products" in st.session_state:
        st.session_state[f"{key_base}_available_products"] = sorted(
            [
                product.name for product in user.products 
                if product.name not in st.session_state[f"{key_base}_disliked_products"]
            ]
        )
        available_products = st.session_state[f"{key_base}_available_products"]
    else:
        available_products = [product.name for product in user.products]
    
    with st.expander("✅ Хочу сегодня съесть", expanded=True):
        remaining_prots = proteins_norm
        remaining_fats = fats_norm
        remaining_carbs = carbs_norm
        for lp in liked_products:
            for p in user.products:
                if lp.name == p.name:
                    remaining_prots -= lp.min_weight * p.proteins_100g / 100
                    remaining_fats -= lp.min_weight * p.fats_100g / 100
                    remaining_carbs -= lp.min_weight * p.carbs_100g / 100
                  
                        
        def __like_product(form_container, remaining_p, remaining_f, remaining_c):
            min_w = st.session_state[f"{key_base}_liked_product_min_weight"]
            max_w = st.session_state[f"{key_base}_liked_product_max_weight"]
            if min_w > max_w:
                with form_container:
                    st.error("Максимум желаемого продукта меньше минимума!", icon="🚫")
                return
            
            product_name = st.session_state[f"{key_base}_liked_product_name"]
            if product_name is None:
                return
            product = [product for product in user.products 
                       if product.name == product_name][0]
            
            max_w_by_prots = remaining_p / (product.proteins_100g / 100) \
                if product.proteins_100g != 0 else float("inf")
            max_w_by_fats = remaining_f / (product.fats_100g / 100) \
                if product.fats_100g != 0 else float("inf")
            max_w_by_carbs = remaining_c / (product.carbs_100g / 100) \
                if product.carbs_100g != 0 else float("inf")
                
            max_possible_weight = min([max_w_by_prots, max_w_by_fats, max_w_by_carbs])
            
            if max_possible_weight < min_w:
                with form_container:
                    st.error(
                        body=f'''
                        Максимум согласно оставшемся БЖУ для:  
                        "{product_name}":  
                        {max_possible_weight:.1f} г  
                        Ваше минимальное значение больше этого числа!
                        ''',
                        icon="🚫",
                    )
                return
            
            liked_product = __LikedProduct(product_name, min_w, max_w)
            if liked_product.name not in [p.name for p in liked_products]:
                liked_products.append(liked_product)
        
        
        st.markdown(
            f"""
            > *Еще возможно использовать:*  
            > - **белков**: {remaining_prots:.2f} г,  
            > - **жиров**: {remaining_fats:.2f} г,  
            > - **углеводов**: {remaining_carbs:.2f} г
            """
        )
        form_container = st.form(key=f"{key_base}_liked_product_form")
        with form_container:
            st.selectbox(
                label="Что хочу",
                options=available_products,
                key=f"{key_base}_liked_product_name",
            )
            st.number_input(
                label="Минимальный вес (гр)",
                min_value=1,
                step=1,
                key=f"{key_base}_liked_product_min_weight",
            )
            st.number_input(
                label="Максимальный вес (гр)",
                min_value=1,
                step=1,
                key=f"{key_base}_liked_product_max_weight",
            )
            st.form_submit_button(
                label="Подтвердить добавление в рацион",
                type="primary",
                on_click=__like_product,
                kwargs={
                    "form_container": form_container,
                    "remaining_p": remaining_prots,
                    "remaining_f": remaining_fats,
                    "remaining_c": remaining_carbs,
                },
            )
            
            
        def __remove_liked_product(product: __LikedProduct):
            liked_products.remove(product)


        for i in range(len(liked_products)):
            with st.container(border=True):
                name_col, min_pfc_col, min_max_w_col, remove_col = st.columns(4)
                lp = liked_products[i]
                product = [p for p in user.products if p.name == lp.name][0]
                with name_col:
                    st.write(product.name)
                with min_pfc_col:
                    min_prot = lp.min_weight * product.proteins_100g / 100
                    min_fat = lp.min_weight * product.fats_100g / 100
                    min_carb = lp.min_weight * product.carbs_100g / 100
                    st.markdown(
                        f"""
                        *мин. б:* **{(min_prot):.2f} г**  
                        *мин. ж:* **{(min_fat):.2f} г**  
                        *мин. у:* **{(min_carb):.2f} г**
                        """
                    )
                with min_max_w_col:
                    st.write(
                        f"""
                        *мин. вес:* **{lp.min_weight:.1f} г**  
                        *макс. вес:* **{lp.max_weight:.1f} г**
                        """
                    )
                with remove_col:
                    st.button(
                        label="Убрать", 
                        key=f"{key_base}_remove_product_{i}_from_liked", 
                        on_click=__remove_liked_product, 
                        args=[lp],
                        use_container_width=True,
                    )
    
    return liked_products


def __show_products(products: list, weights: dict):
    for i in range(len(products)):
        x = weights[i]
        if x != 0:
            with st.container(border=True):
                columns = st.columns(5)
                columns[0].markdown(products[i].name)
                columns[1].markdown(f"вес: {x:.2f} гр")
                columns[2].markdown(
                    f"""
                    белки: {(products[i].proteins_100g * x / 100):.2f} гр  
                    жиры: {(products[i].fats_100g * x / 100):.2f} гр  
                    углеводы: {(products[i].carbs_100g * x / 100):.2f} гр
                    """
                )
                energy = (
                    (products[i].proteins_100g * x / 100 * 4.1) + \
                        (products[i].fats_100g * x / 100 * 9.3) + \
                            (products[i].carbs_100g * x / 100 * 4.1)
                )
                columns[3].markdown(f"калории: {energy:.2f} ккал")
                columns[4].markdown(f"цена: {(products[i].price_100g * x / 100):.2f} руб")


def diet_compiler(key_base: str="default"):
    if "username" not in st.session_state:
        return
    user = SynOrmUserQueries.get_user_by_username(st.session_state["username"])
    
    daily_norm = user.get_daily_norm()
    st.markdown(
        body=f"""
        ## Целевое распределение БЖУ:
        - **белки**:      {daily_norm[0]:.2f} г  
        - **жиры**:       {daily_norm[1]:.2f} г  
        - **углеводы**:   {daily_norm[2]:.2f} г
        """, 
        help="""Итог может незначительно (+-5%) отличаться от цели.  
        **Калорийность** по формуле Миффлина - Сан-Жеора,  
        а конкретные БЖУ - из расчета **1 г белка и 0,8 г жира на 1 кг веса**  
        с увеличением относительно физической активности.  
        Также белок увеличивается на 20%, если цель - похудеть или набрать вес,  
        так как в похудении он участвует в расщеплении жирных кислот,  
        а при наборе веса - в наборе полезной мышечной массы.
        """,
    )
    
    daily_norm = __change_daily_norm(daily_norm, key_base)
    st.plotly_chart(get_pfc_pie_figure(*daily_norm))

    disliked_product_names = __dislike_products(key_base)
    
    liked_products = __like_products(*daily_norm, key_base)
    
    common_max_weight = st.number_input(
        label="Макс. вес других продуктов в гр (нажмите Enter для применения)",
        help="Применяется к тем продуктам, что не указаны выше "
        "и имеют максимальный вес в базе данных больше заданного. "
        "Также вы можете скроллить колесиком мышки, если сделаете поле активным!",
        min_value=20,
        value=100,
        step=10,
        key=f"{key_base}_common_max_weight",
    )
    
    st.divider()
    
    if st.button(
        label="Сотворить рацион!",
        key=f"{key_base}_compile_diet",
        type="primary",
        use_container_width=True,
    ):
        products = [
            product for product in user.products 
            if product.name not in disliked_product_names
        ]
        
        if len(products) == 0:
            st.error("Вы не добавили ни одного продукта!", icon="⛔")
            return
        
        prices_per_g = [product.price_100g / 100 for product in products]
        
        # for -A_ub * x <= -b_ub aka A_ub * x >= b_ub
        lower_bound_coefs = list(
            zip(
                *[
                    [
                        -p.proteins_100g / 100, 
                        -p.fats_100g / 100, 
                        -p.carbs_100g / 100
                    ]
                    for p in products
                ]
            )
        )
        # for A_ub * x <= b_ub
        upper_bound_coefs = []
        for lb_coefs_i in lower_bound_coefs:
            upper_bound_coefs.append(tuple(-coef for coef in list(lb_coefs_i)))
        
        weight_bounds = []
        for p in products:
            liked = False
            for lp in liked_products:
                if p.name == lp.name:
                    liked = True
                    weight_bounds.append((lp.min_weight, lp.max_weight))
            if not liked:
                product_max = p.max_per_day if p.max_per_day else float("inf")
                if common_max_weight < product_max:
                    product_max = common_max_weight
                weight_bounds.append((0, product_max))
                
        # Checking products for compliance with needs
        all_max_weights = [wb[1] for wb in weight_bounds]
        sum_all_products = -np.matrix(lower_bound_coefs) @ np.transpose(np.matrix(all_max_weights))
        is_weights_low = False
        if sum_all_products[0] < daily_norm[0]:
            is_weights_low = True
            st.error(
                body=f"""
                *Недостаточно **белков***!  
                - Вам необходимо: *{daily_norm[0]:.2f} г*,  
                - Сумма белков продуктов согласно ограничению 
                на макс. вес: *{float(sum_all_products[0]):.2f} г*  
                  
                ***Для получения рациона необходимо:***  
                - либо добавить продукт 
                с достаточным содержанием белка (мясо, яйца и др.);  
                - либо увеличить заданные макс. ограничения на вес продуктов, 
                содержащих белок!
                """, 
                icon="⛔",
            )
        if sum_all_products[1] < daily_norm[1]:
            is_weights_low = True
            st.error(
                body=f"""
                *Недостаточно **жиров***!  
                - Вам необходимо: *{daily_norm[1]:.2f} г*,  
                - Сумма жиров продуктов согласно ограничению 
                на макс. вес: *{float(sum_all_products[1]):.2f} г*  
                  
                ***Для получения рациона необходимо:***  
                - либо добавить продукт 
                с достаточным содержанием жира (масло, семена, орехи и др.);  
                - либо увеличить заданные макс. ограничения на вес продуктов, 
                содержащих жир!
                """, 
                icon="⛔",
            )
        if sum_all_products[2] < daily_norm[2]:
            is_weights_low = True
            st.error(
                body=f"""
                *Недостаточно **углеводов***!  
                - Вам необходимо: *{daily_norm[2]:.2f} г*,  
                - Сумма углеводов продуктов согласно ограничению 
                на макс. вес: *{float(sum_all_products[2]):.2f} г*  
                  
                ***Для получения рациона необходимо:***  
                - либо добавить продукт 
                с достаточным содержанием углеводов (фрукты, крупы, сладости и др);  
                - либо увеличить заданные макс. ограничения на вес продуктов, 
                содержащих углеводы!
                """, 
                icon="⛔",
            )
        if is_weights_low:
            return
                
                    
        def get_result_diet(
            prices_per_g, 
            lower_bound_coefs, 
            upper_bound_coefs, 
            daily_norm, 
            upper_coef,
        ):
            return linprog(
                c=prices_per_g, 
                A_ub=[
                    *lower_bound_coefs,
                    *upper_bound_coefs
                ], 
                b_ub=[
                    *[-nutr for nutr in daily_norm],
                    *[nutr*upper_coef for nutr in daily_norm]
                ],
                bounds=weight_bounds,
                integrality=0,
            )
        
        
        default_upper_coef = 1.05
        with st.spinner("Расчет рациона..."):
            result = get_result_diet(
                prices_per_g=prices_per_g,
                lower_bound_coefs=lower_bound_coefs,
                upper_bound_coefs=upper_bound_coefs,
                daily_norm=daily_norm,
                upper_coef=default_upper_coef,
            )
        
        if not result["success"]:
            st.warning(
                body="Текущие продукты и их ограничения по весу не соответствуют "
                "норме по БЖУ. Далее будет произведен расчет ближайшего "
                "возможного рациона с превышением питательных веществ более чем на 5% "
                "(рекомендуемой нормы отклонения)",
                icon="⚠️",
            )
        
        upper_coef = default_upper_coef
        step = 0.01
        while not result["success"]:
            upper_coef += step
            result = get_result_diet(
                prices_per_g=prices_per_g,
                lower_bound_coefs=lower_bound_coefs,
                upper_bound_coefs=upper_bound_coefs,
                daily_norm=daily_norm,
                upper_coef=upper_coef,
            )

            if not result["success"] and (
                round(default_upper_coef, 2) == round(upper_coef - step, 2)
            ): 
                st.write(
                    f"{(upper_coef * 100 - 100):.0f}% превышения нормы: "
                    f"Неудача..."
                )
                st.write("...")
            elif result["success"]:
                st.write(f"{(upper_coef * 100 - 100):.0f}% превышения нормы:")
        
        range_prs = range(len(products))
        diet_proteins = sum(
            [result["x"][i] * products[i].proteins_100g / 100 for i in range_prs]
        )
        diet_fats = sum(
            [result["x"][i] * products[i].fats_100g / 100 for i in range_prs]
        )
        diet_carbs = sum(
            [result["x"][i] * products[i].carbs_100g / 100 for i in range_prs]
        )
        diet_energy = diet_proteins * 4.1 + diet_fats * 9.3 + diet_carbs * 4.1
        diet_price = np.dot(result["x"], prices_per_g)
                
                
        st.subheader("Ваш рацион на день:")        
        if result["success"]:
            __show_products(products=products, weights=result["x"])
            
        st.subheader("Статистика:")
        st.markdown(
            f"""
            > *Общая стоимость: {diet_price:.2f} руб*  
            > *Общая калорийность: {diet_energy:.2f} ккал*  
            > - *Белки: {diet_proteins:.2f} г*  
            > - *Жиры: {diet_fats:.2f} г*  
            > - *Углеводы: {diet_carbs:.2f} г*  
            """
        )
        
        diet_diffs_to_norm = [
            diet_proteins - daily_norm[0], 
            diet_fats - daily_norm[1], 
            diet_carbs - daily_norm[2],
        ]
        pfc_diff = {
            "Питательные вещества": ["Белки", "Жиры", "Углеводы"] * 2,
            "Вес (гр)": [
                *daily_norm,
                *diet_diffs_to_norm
            ],
            "Категория": ["Идеальная норма"] * 3 + ["Избыток диеты"] * 3,
        }
        st.bar_chart(
            data=pfc_diff,
            x="Питательные вещества",
            y="Вес (гр)",
            color="Категория",
        )
            
        if upper_coef > default_upper_coef:
            liked_products.sort(key=lambda lp: lp.name)
            lp_products = []
            for lp in liked_products:
                for p in products:
                    if p.name == lp.name:
                        lp_products.append(p)
                        break
    
            liked_part_of_diet = [
                sum(
                    [
                        liked_products[i].min_weight * lp_products[i].proteins_100g / 100 
                        for i in range(len(lp_products))
                    ]
                ),
                sum(
                    [
                        liked_products[i].min_weight * lp_products[i].fats_100g / 100 
                        for i in range(len(lp_products))
                    ]
                ),
                sum(
                    [
                        liked_products[i].min_weight * lp_products[i].carbs_100g / 100 
                        for i in range(len(lp_products))
                    ]
                )
            ]
            
            # Nutrs that can be changed with linprog (not liked_product min)
            variable_diet_nutrs = [
                diet_proteins - liked_part_of_diet[0],
                diet_fats - liked_part_of_diet[1],
                diet_carbs - liked_part_of_diet[2],
            ]
            
            daily_norm_without_lp = [
                daily_norm[i] - liked_part_of_diet[i] 
                for i in range(len(daily_norm))
            ]
            
            nutrs_coefs_overflow = [variable_diet_nutrs[i] / daily_norm_without_lp[i] 
                                    if daily_norm_without_lp[i] != 0 else 0
                                    for i in range(len(daily_norm_without_lp))]
            
            corrected_diet = [nutr / max(nutrs_coefs_overflow) 
                              for nutr in variable_diet_nutrs]
            
            missing_nutrs = [daily_norm_without_lp[i] - corrected_diet[i] 
                             for i in range(len(daily_norm_without_lp))]
            st.info(
                body=f"""
                Для исправления диеты достаточно **добавить продукты**, 
                либо **изменить существующие максимальные значения продуктов** 
                таким образом, что в ***сумме они дадут при допустимых весах 
                следующее распределение питательных веществ*** (или кратное ему):  
                - Белки: {missing_nutrs[0]:.2f} г  
                - Жиры: {missing_nutrs[1]:.2f} г  
                - Углеводы: {missing_nutrs[2]:.2f} г  
                """,
                icon="‼️",
            )
            
            norm_energy = \
                daily_norm[0] * 4.1 + daily_norm[1] * 9.3 + daily_norm[2] * 4.1
            energy_overflow_to_norm = diet_energy - norm_energy
            st.warning(
                body=f"""Подобный рацион соответствует вашим желаниям и содержит 
                все нужные питательные вещества.  
                ***Но*** он имеет переизбыток в 
                ***{energy_overflow_to_norm:.0f}*** ккал 
                относительно целевой калорийности!
                """,
                icon="⚠️",
            )
            
            metabolism_enerby = norm_energy / user.purpose.coef
            energy_overflow_to_metabolism = diet_energy - metabolism_enerby
            
            human_fat_energy_per_kg = 7700
            fat_gain_velocity_per_month = \
                energy_overflow_to_metabolism / human_fat_energy_per_kg * 30
            if fat_gain_velocity_per_month < 0:
                st.warning(
                    body=f"""С учетом ваших параметров и активности это вызовет:  
                    **меньшую скорость похудения в 
                    {fat_gain_velocity_per_month:.2f} кг в месяц**
                    """,
                    icon="⚠️",
                )
            else:
                st.warning(
                    body=f"""С учетом ваших параметров и активности это вызовет:  
                    **набор веса в 
                    {fat_gain_velocity_per_month:.2f} кг в месяц**
                    """,
                    icon="⚠️",
                )
                    